# utils.py
# 包含辅助函数。

import torch
from tqdm import tqdm
import numpy as np

def get_logits_and_labels(model, data_loader, device):
    """
    辅助函数，用于获取模型在给定数据集上的logits和真实标签。
    返回的tensors保持在指定的device上，而不是移回CPU。
    """
    model.to(device)
    model.eval()
    all_logits = []
    all_labels = []
    with torch.no_grad():
        for inputs, labels in tqdm(data_loader, desc="Extracting Logits/Features"):
            inputs = inputs.to(device)
            logits = model(inputs)
            # 保持在指定设备上，不移回CPU
            all_logits.append(logits)
            all_labels.append(labels.to(device))
            
    return torch.cat(all_logits), torch.cat(all_labels)

def display_results(results, config):
    """
    在控制台打印最终的性能对比结果。
    """
    print("\n\n" + "="*50)
    print(" " * 10 + "FEDERATED ADAPTATION FINAL RESULTS")
    print("="*50)
    print(f"Dataset: {config['dataset']}")
    print(f"Number of Clients: {config['num_clients']}")
    print(f"True Target Priors: {np.round(config['true_target_priors'], 3)}")
    print("-"*50)
    print(f"  - Accuracy (Original Model):      {results['original']:.4f}")
    print(f"  - Accuracy (After Fused Calib.):  {results['calibrated']:.4f}")
    print(f"  - Accuracy (Final Adapted Model): {results['adapted']:.4f}")
    print("="*50)
